import os
from torch.utils.data import Dataset, DataLoader,random_split
import torch
import pickle
from PIL import Image
import torchvision.transforms as transforms
from config import *


tranform_img=transforms.Compose([transforms.Resize((224,224)),
        transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])


class  DataProcess(Dataset):
    def __init__(self,txt_path): # 把数据集的描述文件路径传进来，格式为txt ,每一行是图片的路径和标签
        super(DataProcess, self).__init__()
        with open(txt_path, 'r', encoding='utf-8') as f:
            imgs_info = f.readlines()
        self.imgs_info = list(map(lambda x:x.strip().split('\t'), imgs_info))

    def __len__(self):
        return len(self.imgs_info)

    def __getitem__(self, index):
        image_name, label = self.imgs_info[index]
        image = Image.open(image_name).convert('RGB')
        x= tranform_img(image)
        y = torch.LongTensor([int(label)])
        return x,y

data=DataProcess('data/train.txt')
trainSet,valSet,testSet=random_split(data,lengths=[0.7,0.2,0.1])
trainloader=DataLoader(trainSet,batch_size=batch_size,shuffle=True, num_workers=0, pin_memory=True ) # , num_workers=4, pin_memory=True子进程数并行加载,加快数据加载速度。预先加载
valloader=DataLoader(valSet,batch_size=batch_size,shuffle=True)
testloader=DataLoader(testSet,batch_size=batch_size,shuffle=True)
print("数据处理加载完成")


# for x,y in trainloader:
#     print(x)
#     print(y)



